from argparse import ArgumentParser
from project_utils.image_io import *

import matplotlib.pyplot as plt
import numpy as np
import os.path as osp


def histogram_eq(grayscale):
    """Enhance the contrast and visual quality of an image using histogram equalization.

    `grayscale`: the source grayscale image of int/uint dtype in shape of `(h, w)`.
    """

    levels = np.iinfo(grayscale.dtype).max + 1
    cdf = calculate_image_cdf(grayscale, levels, levels)

    cdf_lut = np.asarray(cdf * (levels - 1), dtype=grayscale.dtype)
    equalized_image = cdf_lut[grayscale]

    return equalized_image


def histogram_match(grayscale, reference):
    """Adjust the intensity of an image to match a reference distribution.

    `grayscale`: the source grayscale image of int/uint dtype in shape of `(h, w)`.

    `reference`: the reference distribution, which can be a target grayscale image of int/uint
    dtype in shape of `(h, w)`, a target histogram in shape of `(bins, )` or a callable that
    are called to generate the pdf of the target intensity I where I falls within [0.0, 1.0).
    """

    if callable(reference):
        aligned_source = align_image_dtype(grayscale, grayscale)
        aligned_dtype = grayscale.dtype
        levels = np.iinfo(aligned_dtype).max + 1
        src_cdf = calculate_image_cdf(aligned_source, levels, levels)
        pdf = reference(np.linspace(0.0, 1.0, levels)[:-1])
        tgt_cdf = calculate_hist_cdf(pdf, 1.0 / levels, levels)

    elif reference.ndim == 1:
        aligned_source = align_image_dtype(grayscale, grayscale)
        aligned_dtype = grayscale.dtype
        levels = np.iinfo(aligned_dtype).max + 1
        src_cdf = calculate_image_cdf(aligned_source, len(reference), levels)
        tgt_cdf = calculate_hist_cdf(reference, levels / len(reference), levels)

    elif reference.ndim == 2:
        aligned_source = align_image_dtype(grayscale, reference)
        aligned_dtype = reference.dtype
        levels = np.iinfo(aligned_dtype).max + 1
        src_cdf = calculate_image_cdf(aligned_source, levels, levels)
        tgt_cdf = calculate_image_cdf(reference, levels, levels)

    else:  # Unexpected reference object
        raise RuntimeError(
            "Expecting a grayscale image, a target histogram or a callable as ref",
            reference,
        )

    cdf_lut = np.interp(src_cdf, tgt_cdf, np.arange(levels))
    cdf_lut = np.asarray(cdf_lut, dtype=aligned_dtype)
    matched_image = cdf_lut[aligned_source]

    return matched_image


def align_image_dtype(source_image, target_image):
    """Convert dtype of source image to that of target image."""

    src_float32 = np.asarray(source_image, dtype=np.float32)
    if not isinstance(source_image.dtype, np.floating):
        src_float32 = src_float32 / np.iinfo(source_image.dtype).max

    if not isinstance(target_image.dtype, np.floating):
        src_float32 = src_float32 * np.iinfo(target_image.dtype).max
    aligned_image = np.asarray(src_float32, dtype=target_image.dtype)

    return aligned_image


def calculate_cdf(coarse_hist, edges, levels):
    """Calculate the cumulative distribution function from coarse histogram, edges and levels."""

    indices = np.searchsorted(edges, np.arange(0, levels, 1), side="right") - 1
    cdf = np.cumsum(coarse_hist[indices])

    return cdf


def calculate_hist_cdf(hist, bin_width, levels):
    """Calculate the cdf from the histogram of the target reference."""

    coarse_hist = hist / np.sum(hist) / bin_width
    edges = np.linspace(0.0, levels, len(hist))
    cdf = calculate_cdf(coarse_hist, edges, levels)

    return cdf


def calculate_image_cdf(image, bins, levels):
    """Calculate the cdf of the intensity for the source image."""

    coarse_hist, edges = np.histogram(image, bins=bins, range=(0, levels), density=True)
    cdf = calculate_cdf(coarse_hist, edges, levels)

    return cdf


def unit_test_problem_01(source_path, target_path):
    """Problem 01: implement histogram equalization algorithm."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    target_image = histogram_eq(source_image)
    save_grayscale_image(target_path, target_image, False)


def unit_test_problem_02(source_path, refimg_path, target_path):
    """Problem 02: implement histogram matching algorithm."""

    source_image = load_image_as_grayscale(source_path, "rgb")
    refimg_image = load_image_as_grayscale(refimg_path, "rgb")
    target_image = histogram_match(source_image, refimg_image)
    save_grayscale_image(target_path, target_image, False)


def plot_histogram(fig, grayscale):
    """Plot intensity histogram of the grayscale image on the figure."""

    levels = np.iinfo(grayscale.dtype).max + 1
    hist, edges = np.histogram(grayscale, levels, (0, levels), density=True)

    plt.figure(fig)
    plt.bar(edges[:-1], hist, 1.2, align="edge", color="dimgray")
    plt.grid(visible=True, which="both", axis="both", linestyle="--")


def equalization_gui(source_rgb_image):
    """Perform histogram equalization to the source RGB image."""

    source_image = convert_to_grayscale_1(source_rgb_image, "rgb")
    target_image = histogram_eq(source_image)

    source_hist = plt.figure("source_eq_hist")
    target_hist = plt.figure("target_eq_hist")

    plot_histogram(source_hist, source_image)
    plot_histogram(target_hist, target_image)

    source_image = convert_to_grayscale_3(source_image)
    target_image = convert_to_grayscale_3(target_image)

    return source_image, source_hist, target_image, target_hist


def matching_gui(ref_type, source_rgb_image, refimg_rgb_image):
    """Perform histogram matching to the source RGB image, given a reference RGB image."""

    assert ref_type in [
        "Image",
        "Histogram",
    ], "Use an image or target histogram for reference."

    source_image = convert_to_grayscale_1(source_rgb_image, "rgb")

    if ref_type == "Image":
        refimg_image = convert_to_grayscale_1(refimg_rgb_image, "rgb")
        target_image = histogram_match(source_image, refimg_image)

    if ref_type == "Histogram":
        target_image = histogram_match(source_image, reference_customized_hist)

    source_hist = plt.figure("source_m_hist")
    target_hist = plt.figure("target_m_hist")
    plot_histogram(source_hist, source_image)
    plot_histogram(target_hist, target_image)
    source_image = convert_to_grayscale_3(source_image)
    target_image = convert_to_grayscale_3(target_image)

    if ref_type == "Image":
        refimg_hist = plt.figure("refimg_m_hist")
        plot_histogram(refimg_hist, refimg_image)
        refimg_image = convert_to_grayscale_3(refimg_image)

    if ref_type == "Histogram":
        refimg_hist, refimg_image = None, None

    return (
        source_image,
        source_hist,
        refimg_image,
        refimg_hist,
        target_image,
        target_hist,
    )


customized_pdf_funcname = "customized_hist"
customized_pdf_template = """
# Numpy has been imported automatically with:
#   `import numpy as np`
# Note that all other packages are unavailable

def {funcname:s}(in_levels: np.ndarray):
  \"""Calculate the histogram base on the input intensity levels, returns the
  number of pixels (int/float) that falls into each level (bin).

  Please DO NOT change the name of this function, which will be invoked
  to calculate the target histogram.

  `in_levels`: a 1D array of input intensity levels in range [0.0, 1.0], whose
  length are determined by the number of intensity levels in the source image.
  \"""

  def gaussian_distribution(x, mu, sig):
    expo = (x - mu) ** 2 / (2 * sig ** 2)
    return np.exp(-expo) / (np.sqrt(2 * np.pi) * sig)

  def gaussian_kernel_helper(x, max_num, mu, sig):
    scale = max_num * np.sqrt(2 * np.pi) * sig
    return scale * gaussian_distribution(x, mu, sig)

  bg = 10 * np.ones(shape=in_levels.shape[0])
  k1 = gaussian_kernel_helper(in_levels, 540, 0.3, 0.04)
  k2 = gaussian_kernel_helper(in_levels, 420, 0.68, 0.04)
  k3 = gaussian_kernel_helper(in_levels, 180, 0.46, 0.12)
  k4 = gaussian_kernel_helper(in_levels, 240, 0.90, 0.06)
  k5 = gaussian_kernel_helper(in_levels, 140, 0.07, 0.15)

  return bg + k1 + k2 + k3 + k4 + k5
""".format(
    funcname=customized_pdf_funcname
)


def customized_histogram(pdf_code):
    """Generate a preview plot for the customized pdf of input intensity levels."""

    _globals, _locals = {"__builtins__": {}, "np": np}, {}
    exec(pdf_code, _globals, _locals)

    in_levels = np.linspace(0.0, 1.0, 256, dtype=np.float32)[:-1]
    return _locals[customized_pdf_funcname](in_levels), in_levels


def draw_customized_histogram(figname, hist, in_levels):
    """Plot intensity histogram on the specified figure `figname`."""

    bar_width = (np.max(in_levels) - np.min(in_levels)) / len(in_levels)
    custom_hist = plt.figure(num=figname)
    plt.bar(in_levels, hist, bar_width, align="edge", color="dimgray")
    plt.grid(visible=True, which="both", axis="both", linestyle="--")

    return custom_hist


def preview_pdf(pdf_code):
    """Preview customized PDF (histogram), which is calculated from user code."""

    hist, in_levels = customized_histogram(pdf_code)
    custom_hist = draw_customized_histogram("custom_hist_pv", hist, in_levels)

    return custom_hist


reference_customized_hist = None


def send_pdf(pdf_code):
    """Save customized PDF (histogram), send to Histogram Matching for use."""

    hist, in_levels = customized_histogram(pdf_code)
    custom_hist = draw_customized_histogram("custom_hist_pv", hist, in_levels)

    global reference_customized_hist
    reference_customized_hist = hist

    return custom_hist


def start_graphical_demo():
    """Start GUI demo for histogram equalization and histogram matching."""

    import gradio as gr

    with gr.Blocks(title="Histogram Processing Demo") as demo:
        gr.Markdown(
            r"""
            ## 直方图均衡和直方图匹配图形化交互Demo
            """
        )

        with gr.Tab("Equalization"):
            upload_eq_1 = gr.Image(label="Source Image", show_label=True, type="numpy")

            button_1 = gr.Button("Process")

            with gr.Row():
                source_eq_img = gr.Image(
                    label="Source Grayscale", show_label=True, type="numpy"
                )
                source_eq_hist = gr.Plot(label="Source Histogram", show_label=True)

            with gr.Row():
                target_eq_img = gr.Image(
                    label="Target Grayscale", show_label=True, type="numpy"
                )
                target_eq_hist = gr.Plot(label="Target Histogram", show_label=True)

            button_1.click(
                equalization_gui,
                inputs=upload_eq_1,
                outputs=[source_eq_img, source_eq_hist, target_eq_img, target_eq_hist],
            )

        with gr.Tab("Matching"):
            with gr.Row():
                with gr.Column():
                    upload_m_1 = gr.Image(
                        label="Source Image", show_label=True, type="numpy"
                    )
                    ref_type = gr.Radio(
                        choices=["Image", "Histogram"], value="Image", show_label=False
                    )
                    button_2 = gr.Button("Process")
                with gr.Column():
                    with gr.Tab("Image"):
                        upload_m_2 = gr.Image(
                            label="Reference Image", show_label=True, type="numpy"
                        )
                    with gr.Tab("Histogram"):
                        upload_m_3 = gr.Plot(
                            label="Reference Histogram", show_label=True
                        )

            with gr.Row():
                source_m_img = gr.Image(
                    label="Source Grayscale", show_label=True, type="numpy"
                )
                refimg_m_img = gr.Image(
                    label="Reference Grayscale", show_label=True, type="numpy"
                )
                target_m_img = gr.Image(
                    label="Target Grayscale", show_label=True, type="numpy"
                )

            with gr.Row():
                source_m_hist = gr.Plot(label="Source Histogram", show_label=True)
                refimg_m_hist = gr.Plot(label="Reference Histogram", show_label=True)
                target_m_hist = gr.Plot(label="Target Histogram", show_label=True)

            button_2.click(
                matching_gui,
                inputs=[ref_type, upload_m_1, upload_m_2],
                outputs=[
                    source_m_img,
                    source_m_hist,
                    refimg_m_img,
                    refimg_m_hist,
                    target_m_img,
                    target_m_hist,
                ],
            )

        with gr.Tab("Custom PDF"):
            gr.Markdown(
                r"""
                ### 自定义目标概率密度函数
                """
            )

            with gr.Row():
                pdf_code = gr.Code(
                    value=customized_pdf_template,
                    language="python",
                    label="Histogram of Input Intensity Levels",
                    show_label=True,
                )
                pdf_plot = gr.Plot(label="PDF Plot", show_label=True)

            with gr.Row():
                button_3 = gr.Button("Preview PDF by Plotting")
                button_4 = gr.Button("Send to Histogram Matching")

            button_3.click(preview_pdf, inputs=pdf_code, outputs=pdf_plot)
            button_4.click(send_pdf, inputs=pdf_code, outputs=upload_m_3)

    demo.launch()


def main_procedure(cmdargs):
    if cmdargs.mode == "equalization":
        source_path = osp.abspath(cmdargs.source)
        target_path = osp.abspath(cmdargs.target)
        unit_test_problem_01(source_path, target_path)

    if cmdargs.mode == "matching":
        source_path = osp.abspath(cmdargs.source)
        refimg_path = osp.abspath(cmdargs.refimg)
        target_path = osp.abspath(cmdargs.target)
        unit_test_problem_02(source_path, refimg_path, target_path)

    if cmdargs.mode == "graphic":
        print("Start GUI demo [using backend: gradio] ...")
        start_graphical_demo()


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Histogram Equalization and Histogram Matching."
    )

    parser.add_argument(
        "mode",
        metavar="MODE",
        choices=["equalization", "matching", "graphic"],
        help="Select processing method: 'equalization', 'matching', 'graphic'.",
    )

    parser.add_argument(
        "--source", metavar="SOURCE", type=str, help="Filepath of the source image."
    )
    parser.add_argument(
        "--target", metavar="TARGET", type=str, help="Filepath of the target image."
    )
    parser.add_argument(
        "--refimg",
        metavar="REFIMG",
        type=str,
        help="Reference image for histogram matching.",
    )

    main_procedure(parser.parse_args())
